import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
from langchain.text_splitter import CharacterTextSplitter
import datetime
############# 在前端显示图像 #################
st.set_page_config(page_title="私人AI翻译助理",
                   page_icon='♾️',
                   layout="centered",  #or wide
                   initial_sidebar_state="expanded",
                   menu_items={
                        'Get Help': 'https://docs.streamlit.io/library/api-reference',
                        'Report a bug': "https://www.extremelycoolapp.com/bug",
                        'About': "一个懂你的AI翻译助理"
                                },
                   )
#
#LOCAL MODEL EN-ZH
#---------------------------------
#  Helsinki-NLP/opus-mt-en-zh
Model_ZH = './models/opus-mt-en-zh/'   #torch
#---------------------------------
### HEADER section
st.header("私人AI翻译助理：帮你把英文翻译成中文")
English = st.text_area("", height=240, key="original",placeholder="请输入或者黏贴英文内容...")
col1, col2, col3 = st.columns([2,5,2])
btn_translate = col2.button("✅ 开始翻译", use_container_width=True, type="primary", key='start')
if btn_translate:
    if English:
        Model_ZH = './models/opus-mt-en-zh/'   #torch
        with st.spinner('AI翻译助理准备中...'):
            st.success(' AI翻译助理开始翻译', icon="🆗")
            # 用于分块的文本分离器函数
            text_splitter = CharacterTextSplitter(
                separator = "\n\n",
                chunk_size = 300,
                chunk_overlap  = 0,
                length_function = len,
            )
            # 将文档分块
            st.success(' 文档块文本...', icon="🆗")
            texts = text_splitter.create_documents([English])
            from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
            # 初始化翻译从英文到中文
            tokenizer_tt0zh = AutoTokenizer.from_pretrained(Model_ZH)
            st.success(' 初始AI语言模型...', icon="🆗")
            model_tt0zh = AutoModelForSeq2SeqLM.from_pretrained(Model_ZH)  #Helsinki-NLP/opus-mt-en-zh  or #Helsinki-NLP/opus-mt-it-zh
            TToZH = pipeline("translation", model=model_tt0zh, tokenizer=tokenizer_tt0zh)
            # 遍历块并连接翻译
            finaltext = ''
            start = datetime.datetime.now()
            print('[bold yellow]翻译进行中...')
            for item in texts:
                line = TToZH(item.page_content)[0]['translation_text']
                finaltext = finaltext+line+'\n'
            stop = datetime.datetime.now()
            elapsed = stop - start
            st.success(f'翻译完成于 {elapsed}', icon="🆗")
            print(f'[bold underline green1] Translation generated in [reverse dodger_blue2]{elapsed}[/reverse dodger_blue2]...')
            st.text_area(label="中文翻译：", value=finaltext, height=350)
            st.markdown(f'翻译完成于： **{elapsed}**')
            st.markdown(f"翻译内容包含单词 {len(English.split(' '))} 个")
    else:
        st.warning("请输入您需要翻译的文本内容！", icon="⚠️")